import {
  Flex,
  Text,
  Box,
  Modal,
  ModalBody,
  ModalCloseButton,
  ModalContent,
  ModalHeader,
  ModalOverlay,
  Spacer,
  useDisclosure,
  SimpleGrid,
} from '@chakra-ui/react';
import IAIIconButton from 'common/components/IAIIconButton';
import { SelectItem } from '@mantine/core';
import { createSelector } from '@reduxjs/toolkit';
import { RootState, stateSelector } from 'app/store/store';
import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
import { loraAdded } from 'features/lora/store/loraSlice';
import { MODEL_TYPE_MAP } from 'features/parameters/types/constants';
import { forEach, map } from 'lodash-es';
import { useCallback, useMemo } from 'react';
import { useGetLoRAModelsQuery } from 'services/api/endpoints/models';
import { useTranslation } from 'react-i18next';
import { FaPlus } from 'react-icons/fa';
import TFLoraItem from './TFLoraItem';
const selector = createSelector(
  stateSelector,
  ({ lora }) => ({
    loras: lora.loras,
  }),
  defaultSelectorOptions
);

const ParamLoRASelect = () => {
  const dispatch = useAppDispatch();
  const { t } = useTranslation();
  const { isOpen, onOpen, onClose } = useDisclosure();
  const { loras } = useAppSelector(selector);
  const { data: loraModels } = useGetLoRAModelsQuery();

  const currentMainModel = useAppSelector(
    (state: RootState) => state.generation.model
  );

  const data = useMemo(() => {
    if (!loraModels) {
      return [];
    }

    const data: SelectItem[] = [];

    forEach(loraModels.entities, (lora, id) => {
      if (!lora || Boolean(id in loras)) {
        return;
      }

      const disabled = currentMainModel?.base_model !== lora.base_model;

      data.push({
        value: id,
        label: lora.model_name,
        disabled,
        group: MODEL_TYPE_MAP[lora.base_model],
        tooltip: disabled
          ? `Incompatible base model: ${lora.base_model}`
          : undefined,
      });
    });

    return data.sort((a, b) => (a.disabled && !b.disabled ? 1 : -1));
  }, [loras, loraModels, currentMainModel?.base_model]);

  const handleAdd = useCallback(
    (v: string | null | undefined) => {
      if (!v) {
        return;
      }
      const loraEntity = loraModels?.entities[v];

      if (!loraEntity) {
        return;
      }

      dispatch(loraAdded(loraEntity));
    },
    [dispatch, loraModels?.entities]
  );

  if (loraModels?.ids.length === 0) {
    return null;
  }

  return (
    <>
      <Flex minWidth="max-content" alignItems="center" py={1}>
        <Box fontSize="sm" fontWeight="600">
          {t('parameters.lora')}
        </Box>
        <Spacer />
        <IAIIconButton
          aria-label={t('parameters.addLora')}
          tooltip={t('parameters.addLora')}
          icon={<FaPlus />}
          size="xs"
          rounded="base"
          bg="base.900"
          _hover={{ bg: 'base.850' }}
          onClick={onOpen}
        />
      </Flex>
      <Modal
        scrollBehavior="inside"
        isOpen={isOpen}
        onClose={onClose}
        size={{ base: 'xl', md: '2xl', lg: '3xl', xl: '4xl', '2xl': '5xl' }}
        isCentered
      >
        <ModalOverlay />
        <ModalContent bg="base.950" fontFamily="Inter">
          <ModalCloseButton />
          <ModalHeader fontWeight="bold" py={2}>
            {t('parameters.lora')}
          </ModalHeader>
          <ModalBody>
            <Flex
              justify="center"
              padding="0 0 1.5rem 0"
              width="100%"
              height="100%"
            >
              <SimpleGrid
                columns={{ base: 2, md: 4, lg: 5, xl: 6 }}
                spacing={4}
              >
                {map(data, (lora) => (
                  <TFLoraItem
                    key={lora.label}
                    {...lora}
                    id={lora.value}
                    onClick={handleAdd}
                  />
                ))}
              </SimpleGrid>
            </Flex>
          </ModalBody>
        </ModalContent>
      </Modal>
    </>
  );
};

export default ParamLoRASelect;
